from pydantic import BaseModel
from fastapi import FastAPI
from . import login
from . import userdb
from .userobj import User
import pickle
import asyncio


# 临时的会话保存
now_session = {}


class WindowsSession(BaseModel):
    windows: list


def load_session(app: FastAPI):
    global now_session

    # 存取会话
    @app.get("/api/session/window")
    @login.check_cookie
    async def get_session_window(user: User):
        global now_session
        if (user.demo):
            return {"windows": []}
        if (user.username in now_session):
            return now_session[user.username]
        ret = await userdb.execute("SELECT `session` FROM `session` WHERE username = ?", False, (user.username,))
        if (ret):
            now_session[user.username] = pickle.loads(ret[0][0])
            return now_session[user.username]
        return {"windows": []}


async def set_session_window(user: User, session_window):  # 设置会话，被 ws 监听器调用
    global now_session
    if (user.demo):
        return
    now_session[user.username] = {
        "windows": session_window["windows"]
    }

    async def save_session():
        await userdb.execute("INSERT OR REPLACE INTO `session` (username, `session`) VALUES (?, ?)",
                             True, (user.username, await asyncio.to_thread(lambda: pickle.dumps(now_session[user.username]))))
    asyncio.create_task(save_session())
